#!/usr/bin/env python3
from utils.global_utils import *
from utils.converter import *
from utils.visualization import *
from event_ros.msg import CameraData, Float64Mat
import pygame

publisherInitialized = False
nodePub = None

subscriberInitialized = False
nodeSub = None

save_dir: str = os.path.join(config["data_dir"].replace("~", os.environ["HOME"]), "camera")
print("Files will be save to %s." % (save_dir,))
H, W, h_btn = config["size"]["h"], config["size"]["w"], 48
is_recording = False
record_btn = None
recorded_events = np.zeros((0, 4), dtype = np.float64)
event_cache = np.zeros((0, 4), dtype = np.float64)


def nodeSubCallback(msg: CameraData) -> None:
    global H, W, h_btn, record_btn, is_recording, event_cache, recorded_events

    if not len(msg.events.data):
        return
    eventData = events_5col_to_4col(np.array(msg.events.data, dtype = np.float64).reshape(-1, len(msg.events.column)))
    if eventData.shape[0]:
        event_cache = np.concatenate([event_cache, eventData], axis = 0)
        if is_recording:
            recorded_events = np.concatenate([recorded_events, eventData], axis = 0)


def subThreadFunc() -> None:
    global subscriberInitialized, nodeSub
    nodeSub = rospy.Subscriber("/celex5_mipi/camera_data", CameraData, nodeSubCallback)
    subscriberInitialized = True
    rospy.spin()


def pubThreadFunc() -> None:
    global publisherInitialized, nodePub, record_btn, is_recording, event_cache, recorded_events
    publisherInitialized = True

    pygame.init()
    screen = pygame.display.set_mode((W, H + h_btn))
    pygame.display.set_caption("Event Viewer")
    font = pygame.font.SysFont('Arial', 24)
    frame_rate = rospy.Rate(config["view"]["fr"])
    while not rospy.is_shutdown():
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
                
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # 左键点击
                    mouse_pos = pygame.mouse.get_pos()
                    
                    # 检查录制按钮点击
                    if record_btn.collidepoint(mouse_pos):
                        if is_recording:
                            events = recorded_events.copy()
                            events[:, 0] -= np.min(events[:, 0])
                            recorded_events = np.zeros((0, 4), dtype = np.float64)
                            if events.shape[0]:
                                now_ts = get_timestamp()
                                now_secs = secs(now_ts)
                                now_nsecs = nsecs(now_ts)
                                sub_dir = os.path.join(save_dir, "events_%d_%d" % (now_secs, now_nsecs))
                                os.makedirs(sub_dir, exist_ok = True)
                                np.save(os.path.join(sub_dir, "events.npy"), events)
                                print("Save %d events to %s." % (events.shape[0], sub_dir))
                            else:
                                print("No event has been recorded.")
                        else:
                            print("Start recording.")
                        is_recording = not is_recording
        
        record_btn = pygame.Rect(W // 4, H, W // 2, h_btn)
        image = white_bg(H, W)
        if event_cache is not None and event_cache.shape[0]:
            image = add_events(image, event_cache)
        event_cache = np.zeros((0, 4), dtype = np.float64)
        image_surface = pygame.surfarray.make_surface(image.swapaxes(0, 1))
        screen.blit(image_surface, (0, 0))
        pygame.draw.rect(screen, (200, 200, 200), pygame.Rect(0, H, W, h_btn))
        btn_color = (0, 200, 0) if not is_recording else (200, 0, 0)
        pygame.draw.rect(screen, btn_color, record_btn)
        btn_text = "Record" if not is_recording else "Stop"
        text_surface = font.render(btn_text, True, (255, 255, 255))
        text = text_surface.get_rect(center = record_btn.center)
        screen.blit(text_surface, text)
        pygame.display.flip()

        frame_rate.sleep()
    
    pygame.quit()


if __name__ == "__main__":
    rospy.init_node("view", anonymous = True)

    subT = threading.Thread(target = subThreadFunc)
    pubT = threading.Thread(target = pubThreadFunc)

    subT.start()
    pubT.start()

    subT.join()
    pubT.join()

    pygame.quit()